package cn.buhler.sink;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.Connection;
import com.rabbitmq.client.ConnectionFactory;
import org.apache.flume.*;
import org.apache.flume.conf.Configurable;
import org.apache.flume.sink.AbstractSink;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class RabbitMQSink extends AbstractSink implements Configurable {

    private static final Logger LOG = LoggerFactory.getLogger(RabbitMQSink.class);

    private String hostname;
    private int port;
    private String username;
    private String password;
    private String queueName;

    private Connection connection;
    private Channel channel;

    private static final int MAX_RETRIES = 3;
    private static final long RETRY_INTERVAL = 5000; // 5 seconds

    @Override
    public void configure(Context context) {
        hostname = context.getString("rabbitmq.hostname", "localhost");
        port = context.getInteger("rabbitmq.port", 5672);
        username = context.getString("rabbitmq.username", "guest");
        password = context.getString("rabbitmq.password", "guest");
        queueName = context.getString("rabbitmq.queueName", "flumeQueue");
    }

    @Override
    public void start() {
        super.start();
        try {
            ConnectionFactory factory = new ConnectionFactory();
            factory.setHost(hostname);
            factory.setPort(port);
            factory.setUsername(username);
            factory.setPassword(password);
            connection = factory.newConnection();
            channel = connection.createChannel();
            channel.queueDeclare(queueName, true, false, false, null);
        } catch (Exception e) {
            LOG.error("Error starting RabbitMQSink", e);
        }
    }

    @Override
    public Status process() throws EventDeliveryException {

        Status status = Status.READY;
        org.apache.flume.Channel channel = getChannel();
        Transaction txn = channel.getTransaction();
        Event event = null;

        int retries = 0;
        boolean isFirstException = true;

        try {
            txn.begin();
            event = channel.take();
            if (event == null) {
                status = Status.BACKOFF;
            } else {
                while (retries < MAX_RETRIES) {
                    try {
                        byte[] messageBody = event.getBody();
                        this.channel.basicPublish("", queueName, null, messageBody);
                        break; // 如果成功，跳出循环
                    } catch (Exception e) {
                        retries++;
                        if (isFirstException || retries == MAX_RETRIES) {
                            LOG.error("Error sending event to RabbitMQ", e);
                            isFirstException = false;
                        }
                        if (retries < MAX_RETRIES) {
                            try {
                                Thread.sleep(RETRY_INTERVAL);
                            } catch (InterruptedException ie) {
                                Thread.currentThread().interrupt();
                            }
                        } else {
                            throw e; // 如果达到最大重试次数，抛出异常
                        }
                    }
                }
            }
            txn.commit();
        } catch (Exception e) {
            txn.rollback();
            status = Status.BACKOFF;
            LOG.error("Failed to process event due to {}", e.getMessage(), e);
        } finally {
            txn.close();
        }

        return status;
    }

    @Override
    public void stop() {
        super.stop();
        try {
            if (channel != null && channel.isOpen()) {
                channel.close();
            }
            if (connection != null && connection.isOpen()) {
                connection.close();
            }
        } catch (Exception e) {
            LOG.error("Error closing RabbitMQ connection", e);
        }
    }
}
